package com.gateway.plugin.resilience4j;

import com.gateway.common.constant.Constants;
import com.gateway.common.dto.RuleData;
import com.gateway.common.dto.SelectorData;
import com.gateway.common.dto.convert.rule.Resilience4JHandle;
import com.gateway.common.enums.PluginEnum;
import com.gateway.common.utils.UriUtils;
import com.gateway.plugin.api.GatewayPluginChain;
import com.gateway.plugin.api.context.GatewayContext;
import com.gateway.plugin.base.AbstractGatewayPlugin;
import com.gateway.plugin.base.util.CacheKeyUtils;
import com.gateway.plugin.resilience4j.builder.Resilience4JBuilder;
import com.gateway.plugin.resilience4j.config.Resilience4JConfig;
import com.gateway.plugin.resilience4j.executor.CircuitBreakerExecutor;
import com.gateway.plugin.resilience4j.executor.CombinedExecutor;
import com.gateway.plugin.resilience4j.executor.Executor;
import com.gateway.plugin.resilience4j.executor.RateLimiterExecutor;
import com.gateway.plugin.resilience4j.handler.Resilience4JHandler;
import org.springframework.http.HttpStatus;
import org.springframework.http.HttpStatusCode;
import org.springframework.web.client.HttpStatusCodeException;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.io.Serial;
import java.util.Objects;
import java.util.function.Function;

/**
 * @program: api-gateway
 * @description:
 * @author: YuKai Fan
 * @create: 2025/2/24 09:39
 **/
public class Resilience4JPlugin extends AbstractGatewayPlugin {

    private final CombinedExecutor combinedExecutor;

    private final RateLimiterExecutor rateLimiterExecutor;

    private final CircuitBreakerExecutor circuitBreakerExecutor;

    public Resilience4JPlugin(CombinedExecutor combinedExecutor, RateLimiterExecutor rateLimiterExecutor, CircuitBreakerExecutor circuitBreakerExecutor) {
        this.combinedExecutor = combinedExecutor;
        this.rateLimiterExecutor = rateLimiterExecutor;
        this.circuitBreakerExecutor = circuitBreakerExecutor;
    }

    @Override
    protected Mono<Void> doExecute(ServerWebExchange exchange, GatewayPluginChain chain, SelectorData selector, RuleData rule) {
        final GatewayContext gatewayContext = exchange.getAttribute(Constants.CONTEXT);
        Objects.requireNonNull(gatewayContext);
        Resilience4JHandle resilience4JHandle = Resilience4JHandler.CACHED_HANDLE.get().obtainHandle(CacheKeyUtils.INST.getKey(rule));
        resilience4JHandle.checkData(resilience4JHandle);
        if (isCircuitBreaker(resilience4JHandle)) {
            return circuitBreaker(exchange, chain, rule);
        }
        if (isRateLimit(resilience4JHandle)) {
            return rateLimiter(exchange, chain, rule);
        }
        if (isCombined(resilience4JHandle)) {
            return combined(exchange, chain, rule);
        }
        return chain.execute(exchange);
    }

    private Mono<Void> rateLimiter(final ServerWebExchange exchange, final GatewayPluginChain chain, final RuleData ruleData) {
        return rateLimiterExecutor.run(
                chain.execute(exchange), fallback(rateLimiterExecutor, exchange, null), Resilience4JBuilder.build(ruleData))
                .onErrorResume(throwable -> rateLimiterExecutor.withoutFallback(exchange, throwable));
    }

    private Mono<Void> circuitBreaker(final ServerWebExchange exchange, final GatewayPluginChain chain, final RuleData ruleData) {
        return circuitBreakerExecutor.run(
                        chain.execute(exchange), fallback(rateLimiterExecutor, exchange, null), Resilience4JBuilder.build(ruleData))
                .onErrorResume(throwable -> rateLimiterExecutor.withoutFallback(exchange, throwable));
    }

    private Mono<Void> combined(final ServerWebExchange exchange, final GatewayPluginChain chain, final RuleData ruleData) {
        Resilience4JConfig config = Resilience4JBuilder.build(ruleData);
        return combinedExecutor.run(
                chain.execute(exchange).doOnSuccess(v -> {
                    HttpStatusCode statusCode = exchange.getResponse().getStatusCode();
                    if (Objects.isNull(statusCode) || !statusCode.is2xxSuccessful()) {
                        exchange.getResponse().setStatusCode(null);
                        throw new CircuitBreakerStatusCodeException(Objects.isNull(statusCode) ? HttpStatus.INTERNAL_SERVER_ERROR : statusCode);
                    }
                }), fallback(combinedExecutor, exchange, config.getFallBackUri()), config);
    }

    private Function<Throwable, Mono<Void>> fallback(final Executor executor,
                                                     final ServerWebExchange exchange,
                                                     final String uri) {
        return throwable -> executor.fallback(exchange, UriUtils.createUri(uri), throwable);
    }

    @Override
    public int getOrder() {
        return PluginEnum.RESILIENCE4J.getCode();
    }

    public boolean isCombined(Resilience4JHandle handle) {
        return handle.getRateLimitEnable() == 1 && handle.getCircuitEnable() == 1;
    }

    public boolean isRateLimit(Resilience4JHandle handle) {
        return handle.getRateLimitEnable() == 1 && handle.getCircuitEnable() == 0;
    }

    public boolean isCircuitBreaker(Resilience4JHandle handle) {
        return handle.getRateLimitEnable() == 0 && handle.getCircuitEnable() == 1;
    }

    public static class CircuitBreakerStatusCodeException extends HttpStatusCodeException {

        @Serial
        private static final long serialVersionUID = 8660297999616276146L;

        public CircuitBreakerStatusCodeException(final HttpStatusCode statusCode) {
            super(statusCode);
        }
    }

}
